import gradio as gr
import pickle
from PIL import Image
import numpy as np
from sklearn.preprocessing import LabelEncoder

# 加载模型
with open('best_model.pkl', 'rb') as f:
    model = pickle.load(f)

# 预处理函数
def preprocess(image):
    # 将PIL图像转换为NumPy数组
    image_array = np.array(image) / 255.0
    # 假设模型需要的输入是64x64的图像
    image_array = image_array.reshape(32, 32)
    return image_array

# 预测函数
def predict(image):
    # 预处理图像
    image_array = preprocess(image)
    # 预测
    prediction = model.predict([image_array])
    # 将预测结果转换为标签
    labels = ['cat', 'dog']
    predicted_label = labels[prediction[0]]
    return predicted_label

# 创建Gradio接口
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),
    title="Cat vs Dog Classifier",
    description="Upload an image to classify it as a cat or a dog."
)

# 启动应用
if __name__ == "__main__":
    iface.launch()